#include "bsd_socket.h"
#include "mbedtinydtls.h"

// tinydtls globals 
dtls_context_t *dtls_context = NULL;
session_t g_dst;
uint8_t *clearText;
int clearTextLen;

static unsigned char psk_id[PSK_ID_MAXLEN] = PSK_DEFAULT_IDENTITY;
static size_t psk_id_length = sizeof(PSK_DEFAULT_IDENTITY) - 1;
static unsigned char psk_key[PSK_MAXLEN] = PSK_DEFAULT_KEY;
static size_t psk_key_length = sizeof(PSK_DEFAULT_KEY) - 1;

/* This function is the "key store" for tinyDTLS. It is called to
 * retrieve a key for the given identity within this particular
 * session. */
static int get_psk_info(struct dtls_context_t *ctx,
        const session_t *session,
        dtls_credentials_type_t type,
        const unsigned char *id, size_t id_len,
        unsigned char *result, size_t result_length) {

//  TRACELOG(printf("\rgetpskinfo\n"));
  switch (type) {
  case DTLS_PSK_IDENTITY:

    if (result_length < psk_id_length) {
      printf("\rcannot set psk_identity -- buffer too small\n");
      return dtls_alert_fatal_create(DTLS_ALERT_INTERNAL_ERROR);
    }

    memcpy(result, psk_id, psk_id_length);
    return psk_id_length;
  case DTLS_PSK_KEY:
    if (id_len != psk_id_length || memcmp(psk_id, id, id_len) != 0) {
      printf("\rPSK for unknown id requested, exiting\n");
      return dtls_alert_fatal_create(DTLS_ALERT_ILLEGAL_PARAMETER);
    } else if (result_length < psk_key_length) {
      printf("\rcannot set psk -- buffer too small\n");
      return dtls_alert_fatal_create(DTLS_ALERT_INTERNAL_ERROR);
    }

    memcpy(result, psk_key, psk_key_length);
    return psk_key_length;
  default:
    printf("\runsupported request type: %d\n", type);
  }

  return dtls_alert_fatal_create(DTLS_ALERT_INTERNAL_ERROR);
}

// this is a required callback for tinydtls it is called
// whenever tinydtls handles a raw buffer through dtls_handle_message.
// tinydtls extracts the unencrypted data and passes it onto this function
static int read_from_peer(struct dtls_context_t *ctx, session_t *session, uint8 *data, size_t len) {
	clearText = data;
	clearTextLen = len;

	return 0;
}

/// this is a required callback for tinydtls it is called
// whenever tinydtls needs to send a raw buffer (i.e whenever dtls_write is called by the application)
static int send_to_peer(struct dtls_context_t *ctx, session_t *session, uint8 *data, size_t len) {
	int fd = *(int *)dtls_get_app_data(ctx);
	return sendto(fd, data, len, MSG_DONTWAIT, &session->addr.sa, session->size);
}

static int dtls_handle_read(struct dtls_context_t *ctx) {
  int fd;
  session_t session;
  #define MAX_READ_BUF 256
  static uint8 buf[MAX_READ_BUF];
  int len;

  fd = *(int *)dtls_get_app_data(ctx);

  if (fd < 0)
    return -1;

  memset(&session, 0, sizeof(session_t));
  session.size = sizeof(session.addr);
  len = recvfrom(fd, buf, MAX_READ_BUF, 0, &session.addr.sa, &session.size);

  if (len < 0) {
	printf("\rGot nothing from read\n");
    perror("recvfrom");
    return -1;
  } else {
    dtls_dsrv_log_addr(DTLS_LOG_DEBUG, "peer", &session);
    dtls_debug_dump("bytes from peer", buf, len);
  }

  return dtls_handle_message(ctx, &session, buf, len);
}

static dtls_handler_t cb = {
  .write = send_to_peer,
  .read  = read_from_peer,
  .event = NULL,
#ifdef DTLS_PSK
  .get_psk_info = get_psk_info,
#endif /* DTLS_PSK */
#ifdef DTLS_ECC
  .get_ecdsa_key = get_ecdsa_key,
  .verify_ecdsa_key = verify_ecdsa_key
#endif /* DTLS_ECC */
};

int tinydtls_init(int *sockfd) {
	int result = 0, counter = 0;
	log_t log_level = DTLS_LOG_WARN;
    // dtls stuff
	dtls_peer_t *peer = NULL;
	fd_set rfds;

	int timeoutNum = 3;

    // dtls init must always be called for memory allocation
    dtls_init();
    dtls_set_log_level(log_level);

dtlsConnectAuth:
    // setup DTLS context
	TRACELOG(printf("\rCreating DTLS context\n"));
    dtls_context = dtls_new_context(sockfd);
    if(!dtls_context) {
        dsrv_log(DTLS_LOG_EMERG, (char *)"cannot create context");
        return -1;
    }

    // forced to use this call back system
    dtls_set_handler(dtls_context, &cb);

    TRACELOG(printf("\rIssuing dtls_connect\n"));
	result = dtls_connect(dtls_context, &g_dst);
    if(result < 0) {
        printf("\rError in dtls_connect: %d", result);
        return -2;
    }
    if(result == 0) {
        printf("\rChannel already exists");
        return -3;
    }
    TRACELOG(printf("\rdtls_connect successfull\n"));

	peer = dtls_get_peer(dtls_context, &g_dst);
	while (true) {
		FD_ZERO(&rfds);
		FD_SET(*sockfd, &rfds);
		struct timeval timeout;
		timeout.tv_sec = 3;
		timeout.tv_usec = 0;

		result = select(*sockfd+1, &rfds, NULL, 0, &timeout);

		if (result > 0) {	/* OK */
			if (FD_ISSET(*sockfd, &rfds))
				dtls_handle_read(dtls_context);
			if (peer->state == DTLS_STATE_CONNECTED) {
				printf("\rPeer connect is OK......................\n");
				break;
			}
		} else if (result == 0) {	/* timeout */
			perror("\rSelect the sockfd timeout, retry to connect the peer...");
			dtls_free_context(dtls_context);
			peer = NULL;
			if (timeoutNum-- == 0)
				return -4;
			goto dtlsConnectAuth;
		} else {	/* error */
			perror("\rSelect error");
		}
	}

	int timeout = 6000;		//timeout 6s
	setsockopt(*sockfd, SOL_SOCKET, SO_SNDTIMEO, (char *)&timeout, sizeof(timeout));
	setsockopt(*sockfd, SOL_SOCKET, SO_RCVTIMEO, (char *)&timeout, sizeof(timeout));

	return 0;
}
